#!/usr/bin/env python
# coding: utf-8
# 二进制遗传算法，寻找Ackley函数的最小值

import numpy as np
import matplotlib.pyplot as plt
import itertools
import random
import math


class Individual:
    def __init__(self, genetype=None, phenotype=None, fitness=None):
        self.genetype = genetype
        self.phenotype = phenotype
        self.fitness = fitness

    def __str__(self):
        return '{0},{1},{2}'.format(self.genetype, self.phenotype, self.fitness)


# phenotype
def phenotype(genetype):
    x = int(genetype[:6], 2)
    y = int(genetype[6:], 2)
    #     print(x,y)
    return -5 + 10 * x / 63, -5 + 10 * y / 63


def ackley(x, y):
    return math.e - 20 * math.exp(-0.2 * math.sqrt((x * x + y * y) / 2)) - math.exp(
        (math.cos(2 * math.pi * x) + math.cos(2 * math.pi * y)) / 2)


def fitness(pop):
    for indv in pop:
        value = -ackley(*indv.phenotype)
        assert value >= 0
        indv.fitness = value


def selection_one(pop, _sum):
    rand = random.random() * _sum
    f = pop[0].fitness
    k = 0
    while f < rand:
        k += 1
        f += pop[k].fitness
    return k


# 选出一对父代
def selection(pop):
    _sum = 0
    for indv in pop:
        _sum += indv.fitness
    j = selection_one(pop, _sum)
    k = selection_one(pop, _sum)
    # # 不能自己交配
    # while k == j:
    #     k = selection_one(pop, _sum)

    return pop[j], pop[k]


def cross(parents):
    # 保证有交叉，限制范围
    rand = random.randrange(1, len(parents[0].genetype) - 1)
    # print(parents[0].genetype, parents[1].genetype)
    child1_gene = parents[0].genetype[:rand] + parents[1].genetype[rand:]
    child2_gene = parents[1].genetype[:rand] + parents[0].genetype[rand:]
    # print(child1_gene, child2_gene)
    child1 = Individual(child1_gene, phenotype(child1_gene))
    child2 = Individual(child2_gene, phenotype(child2_gene))
    return child1, child2


def mutate(children_pop):
    for child in children_pop:
        new_genetype = []
        for bit in child.genetype:
            rand = random.random()
            if rand < mutate_rate: # mutate
                if bit == '1':
                    new_genetype.append('0')
                else:
                    new_genetype.append('1')
            else:
                new_genetype.append(bit)
        child.genetype = ''.join(new_genetype)


def visualize(pop, total_count, count):
    row = total_count/2
    plt.subplot(row, 2, count)
    plt.xlim(-5, 5)
    plt.ylim(-5, 5)
    points = [p.phenotype for p in pop]
    plt.scatter(*zip(*points))


x = ['{:06b}'.format(i) for i in range(64)]
y = ['{:06b}'.format(i) for i in range(64)]
pop_set = [''.join(i) for i in itertools.product(x, y)]

pop_size = 20
mutate_rate = 0.02
parents_pop = [Individual() for i in range(pop_size)]

# initialize - genetype
genetypes = random.choices(pop_set, k=pop_size)
for i in range(pop_size):
    parents_pop[i].genetype = genetypes[i]
    parents_pop[i].phenotype = phenotype(genetypes[i])
print(parents_pop[0])

generation = 20
count = 0

while count < generation:
    # 计算种群中每个父代的适应度
    fitness(parents_pop)
    visualize(parents_pop, generation, count + 1)
    fitness_values = [p.fitness for p in parents_pop]
    print('gen: {0}, best: {1}, avg: {2}'.format(count, max(fitness_values), sum(fitness_values)/pop_size))
    children_pop = []
    while len(children_pop) < len(parents_pop):
        # 用适应度根据概率选出一对交配的父代
        parents = selection(parents_pop)
        # 交配生成子代
        children = cross(parents)
        children_pop += children
    # 变异
    mutate(children_pop)

    parents_pop = children_pop
    count += 1

plt.show()
